'''
Created on Jun 13, 2010

@author: oabalbin
'''
import sys
import numpy as np
#import MySQLdb;
import mysql.connector
from collections import deque, defaultdict

from interactome.common.classes import interaction

class query():
    
    def __init__(self, dbhost, dbuser, dbpasswd, database):
        """ Database Information """
        self.myhost = dbhost
        self.myuser = dbuser
        self.mypasswd = dbpasswd
        self.mydb = database
        self.create_connector()
    
    def create_connector(self):
        #self.conn=MySQLdb.connect(host=self.myhost,user=self.myuser,passwd=self.mypasswd,db=self.mydb)
        self.conn=mysql.connector.Connect(host=self.myhost,user=self.myuser,passwd=self.mypasswd,db=self.mydb)
    
    def close_connector(self, connector):
        connector.close
    
    def get_all_knownGenes_from_ucsc(self):
        
        genes = deque()
        cursor=self.conn.cursor()
        query = 'select distinct geneSymbol, refseq from kgXref;'
        
        cursor.execute(query)
        # Fetch all the rows in a list of lists.
        results = cursor.fetchall()
        for row in results:
            if not row[1]:
                continue
            
            genes.append(row[0])
        #[genes.append(row[0]) for row in results]
        
        return genes
    
    def map_genes2proteins_in_string(self,genelist):
        """
        It assumes that list of genes is in hugo name. It could also receive refseq (transcript or protein), ensembl, agilent or illumina probes.
        """
        #homo_sapiens_idsp='9606.'
        items = defaultdict()
        items_reverse=defaultdict(list)
        notfound = []
        cursor=self.conn.cursor()
        item_names = np.array(['prot','gen'])
        #genes = ['9606.ENSP00000344798','9606.ENSP00000269571']
        
        for gen in genelist:
            query ='select distinct protein_id from protein_aliases where alias=\"' + gen +'\" and source="BLAST_UniProt_GN BioMart_HUGO Ensembl_EntrezGene Ensembl_HUGO Ensembl_UniProt_GN";'
            cursor.execute(query)
            results = cursor.fetchall()
                        
            if not results:
                prot = self.map_genes2proteins_in_string2(gen)
                if not prot:
                    notfound.append(gen)
                    continue    
            else:
                prot = results[0][0] #homo_sapiens_idsp+results[0][0]
            
            
            items[gen] = prot
            items_reverse[prot] = gen
            tmp = np.array([prot,gen])
            item_names = np.vstack((item_names,tmp))

        #print "Not protein equivalent was found for these genes=", len(notfound), notfound      
        
        return items, items_reverse, notfound, item_names[1:,:]
    
    
    def map_genes2proteins_in_string2(self,gen):
        """
        It assumes that list of genes is in hugo name. It could also receive refseq (transcript or protein), ensembl, agilent or illumina probes.
        """
        cursor=self.conn.cursor()        
        query ='select distinct protein_id from protein_aliases where alias=\"' + gen +'\";'
        cursor.execute(query)
        results = cursor.fetchall()
                        
        if not results:
            prot=''
        else:
            prot = results[0][0]
            
        return prot
                
    
    def extract_protein_interactions(self, list_of_proteins):
        """
        It extracts protein interactions for the list of proteins from STRING database
        """
        list_of_interactions=[]
        notfound = []
        cursor=self.conn.cursor()
        for prot in list_of_proteins:
            query = 'select * from protein_links_detailed where item_id_a=\"' + prot +'\";'
            cursor.execute(query)
            rows = cursor.fetchall()
            if not rows:
                notfound.append(prot)
            else:
                for results in rows:
                    item_id_a, item_id_b, neighborhood, fusion, cooccurence, coexpression, experimental, indatabase, textmining, combined_score = \
                    results[0], results[1], int(results[2]), int(results[3]),int(results[4]), int(results[5]), int(results[6]), int(results[7]), int(results[8]), int(results[9])
                    # Do not include self loops in the interactome
                    if item_id_a != item_id_b:    
                        list_of_interactions.append(interaction(item_id_a, item_id_b, neighborhood, fusion, cooccurence, 
                                                                coexpression, experimental, indatabase, textmining, combined_score))
        return  list_of_interactions
    
    def extract_protein_interactions_string8(self, list_of_proteins,trans2prot_edge):
        """
        It extracts protein interactions for the list of proteins from STRING database
        """
        max_edge_score = 1000
        list_of_interactions=[]
        notfound = []
        cursor=self.conn.cursor()
        for item_id_a_hugo, prot in list_of_proteins.iteritems():
            query = 'select A.*, B.alias '+\
                     'from protein_links_detailed as A, protein_aliases as B '+\
                    'where A.item_id_a=\"' + prot +'\" and B.protein_id = A.item_id_b and B.source=\"BLAST_UniProt_GN BioMart_HUGO Ensembl_EntrezGene Ensembl_HUGO Ensembl_UniProt_GN\" order by A.combined_score;'
            
            cursor.execute(query)
            rows = cursor.fetchall()
            
            if not rows:
                notfound.append((prot,item_id_a_hugo))
            else:
                for results in rows:
                    item_id_a, item_id_b, neighborhood, fusion, cooccurence, coexpression, experimental, indatabase, textmining, \
                    combined_score, item_id_b_hugo  = \
                    results[0], results[1], int(results[2]), int(results[3]),int(results[4]), int(results[5]), int(results[6]), \
                    int(results[7]), int(results[8]), int(results[9]),results[10]  
                    # Do not include self loops in the interactome
                    if item_id_a != item_id_b:    
                        list_of_interactions.append(interaction(item_id_a, item_id_b, neighborhood, fusion, cooccurence, 
                                                                coexpression, experimental, indatabase, textmining, combined_score,
                                                                item_id_a_hugo, item_id_b_hugo))
            
            # This introduces the link between the transcript and its protein. 
            # To start with the weight edge between them is maximum expected
            # Therefore, there will be not cost associated  with adding this edge
            # during the pcst algorithm

            if trans2prot_edge:
                item_id_a, item_id_b, neighborhood, fusion, cooccurence, coexpression, experimental, indatabase, textmining, \
                combined_score  = \
                prot+'_mrna', prot, max_edge_score, max_edge_score,max_edge_score, max_edge_score, max_edge_score, \
                max_edge_score, max_edge_score, max_edge_score  

                list_of_interactions.append(interaction(item_id_a, item_id_b, neighborhood, fusion, cooccurence, 
                                                                coexpression, experimental, indatabase, textmining, combined_score,
                                                                item_id_a_hugo+'_mrna', item_id_a_hugo))
                print trans2prot_edge, item_id_a_hugo+'_mrna', item_id_a_hugo
        
        if len(notfound) >0 :
            print "There were not interactions found for the following proteins = %d"%len(notfound)
            print notfound
        
        return  list_of_interactions

    
    def extract_protein_interaction_mode(self, item_a, item_b):
        """
        It returns the mode of the protein interaction if it is reported in the STRING database. 
        """
        cursor=self.conn.cursor()
        query = ' select * from protein_actions_detailed where item_id_a= \"' + item_a +'\" and  item_id_b= \"' + item_b +'\";'
        cursor.execute(query)
        rows = cursor.fetchall()
        if not rows:
            mode = ["unknown", "unknown", 0]
        else:
            mode = [rows[0][2], rows[0][3], rows[0][4]]
        
        return mode
     
    def translate_protname2hugoname(self,protlist):
                    
        genes = deque()
        cursor=self.conn.cursor()
        
        for prot0 in protlist:
            prot = prot0.split('.')[0]
            query = 'select distinct geneSymbol from kgXref where protAcc=\"' + prot +'\";'
            
            cursor.execute(query)
            # Fetch all the rows in a list of lists.
            results = cursor.fetchall()
            if not results:
                genes.append(("notfound", prot0))
                continue    
            genes.append((results[0][0], prot0))
            #[genes.append(row[0]) for row in results]
        
        return genes
        
    def extract_single_p2pinteraction(self, prot_a, prot_b, gen_a, gen_b):
        """
        It extracts protein interactions for the list of proteins from STRING database
        select * from protein_links_detailed where item_id_b='9606.ENSP00000256078' and item_id_a='9606.ENSP00000408695' LIMIT 10;
        prot_a, prot_b ensembl protein names
        gen_a, gen_b hugo gene names
        """

        notfound = []
        cursor=self.conn.cursor()
        query = 'select * from protein_links_detailed where item_id_a=\"' + prot_a +'\" AND item_id_b=\"' + prot_b +'\";'
        cursor.execute(query)
        rows = cursor.fetchall()
        
         
        
        if not rows:
            notfound.append((prot_a,prot_b))
            edge=None
            print 'Interaction was not found in Ensembl=',gen_a, gen_b
        else:
            for results in rows:
                
                item_id_a, item_id_b, neighborhood, fusion, cooccurence, coexpression, experimental, indatabase, textmining, combined_score = \
                results[0], results[1], int(results[2]), int(results[3]),int(results[4]), int(results[5]), int(results[6]), int(results[7]), int(results[8]), int(results[9])
                # Do not include self loops in the interactome
                if item_id_a != item_id_b:    
                    edge=interaction(item_id_a, item_id_b, neighborhood, fusion, cooccurence, 
                                                            coexpression, experimental, indatabase, textmining, combined_score,gen_a, gen_b)
                if gen_a == 'KRAS' or gen_b=='KRAS':
                    print item_id_a, item_id_b, gen_a, gen_b
                    print rows
                    
        return  edge


   
    

            

            
        






